import numpy
from common import *
from dolfin import Cell
from common import directional_diameter

def get_sigma(v,eps):
    '''Stabilization parameter for crosswind method of Johnson(1987).'''
    mesh = v.function_space().mesh()
    nc = mesh.num_cells()
    taus = numpy.zeros(nc)

    # solve problem vMag*phi = sqrt(inner(v,v)) to get the magnitude of velocity
    V = FunctionSpace(mesh,'CG',v.element().degree())
    vNorm = TrialFunction(V)
    vNormTest = TestFunction(V)

    a = vNorm*vNormTest*dx 
    l = sqrt(inner(v,v))*vNormTest*dx

    vNorm = Function(V)
    A = assemble(a)
    b = assemble(l)

    solve(A,vNorm.vector(),b)

    for i in xrange(nc):
        c = Cell(mesh,0)
        x = c.midpoint() 
        vMag = vNorm(x.x(),x.y())
        
        if vMag > 1E-16:
            h = directional_diameter(c,v)
            tau = max(0, vMag*h**(3./2)-eps) # as found in A computational comparison of methods diminishing spurious 
                                             # oscillations in FE solutions of convection-diffusion equations
        else:
            tau = 0
            
        taus[i] = tau
    return taus
    
